# -*- coding: utf-8 -*-

import os, sys

#------------------------------------------------------------------------

class MeshInfo():
  u"""A simple class containing basic mesh information """
  def __init__(self) :
    self.nnodes = 0
    self.ncells = 0
    self.nsides = 0
    self.marked = 0
  def __init__(self, meshname):
    filehandle = open(meshname + u'.pylotmesh/name', u'r')
    toto = filehandle.readline()
    toto = filehandle.readline()
    filehandle.close()
    totosplit = toto.split(u'_')
    self.nnodes = int(totosplit[1])
    self.ncells = int(totosplit[2])
    self.nsides = int(totosplit[3])
    try :
      filehandle = open(meshname + u'.pylotmesh/marked_cells', u'r')
      toto = filehandle.readline()
      totosplit = toto.split()
      self.marked = int(totosplit[0])
    except :
      self.marked = -1


#------------------------------------------------------------------------

class IterationInfo() :
  def __init__(self, rundir, model) :
    self.postprocess = {}
    filename = rundir + u'/NonLinearInfo.pylotsolution'
    try:
      filehandle = open(filename, u'r')
      toto = filehandle.readline()
      totolist = toto.split()
      if len(totolist) != 2:
        print u'$$$ IterationInfo() problem in file \"', filename, u'\" (expecting two ints) I have read ', toto
        sys.exit(1)
      self.nonlinear_niter = totolist[0]
      self.linear_niter = totolist[1]
      filehandle.close()
    except IOError:
      print u'$$$ IterationInfo() file \"', filename, u'\" could not be opened !!!'
    # model = rundir.split('/')[-3]
    filename = rundir + u'/' + model + u'_PostProcess.pylotsolution'
    filehandle = open(filename, u'r')
    modeladj = model + u'_Adjoint'
    filenameadj = rundir + u'/' + modeladj + u'_PostProcess.pylotsolution'
    try:
      toto = filehandle.readlines()
      for line in toto:
        linelist = line.split()
        if len(linelist) < 2:
          print u'$$$ IterationInfo() file \"', filename, u'\" could not be read properly !!!'
          sys.exit(1)
        self.postprocess[linelist[0]] = linelist[1:]
      if filenameadj.exists():
        filehandleadj = open(filenameadj, u'r')
        totoadj = filehandleadj.readlines()
        for line in totoadj:
          linelist = line.split()
          if len(linelist) < 2:
            print u'$$$ IterationInfo() file \"', filename, u'\" could not be read properly !!!'
            sys.exit(1)
          self.postprocess[linelist[0]] = linelist[1:]
        filehandleadj.close()
    except IOError:
      print u'### Could not read file', self.rundir + u'/ErrorEstimator'
      sys.exit(1)
    filehandle.close()

#------------------------------------------------------------------------

class StdAmrLoopVisitor() :
  def __init__(self, pathbook, manager, rundir) :
    self.pathbook = pathbook
    self.manager = manager
    self.rundir = rundir

  def getResultsFileName(self, marker, marking_percentage, gamma):
    return self.rundir + u'/result_' + marker + u'_' + marking_percentage + u'_' + gamma

  def write(self, filehandle, iteration, meshinfo, iterationinfo):
    if iteration == 0:
      outformat = u"%3s %8s %8s %5s %5s     " % (u'iter', u'N', u'M', u'nlit', u'linit')
      for item in iterationinfo.postprocess:
        first = 1
        for subitem in iterationinfo.postprocess[item]:
          if first:
            outformat += u" %12s" % item[-12:]
            first = 0
          else:
            outformat += u"%12s" % u' '
      print outformat
      filehandle.write(u"# " + outformat + u'\n')
    out = u"%3d %8d %8d %5d %5d     " % (int(iteration), meshinfo.ncells, meshinfo.marked, int(iterationinfo.nonlinear_niter), int(iterationinfo.linear_niter))
    for item in iterationinfo.postprocess.keys():
      for subitem in iterationinfo.postprocess[item]:
       out += u" %12.3e" % float(subitem)
    print out
    filehandle.write(out + u'\n')

  def solverPost(self, iter):
    pass

#------------------------------------------------------------------------

class TestVisitor(StdAmrLoopVisitor):
  def __init__(self, pathbook, manager, rundir) :
    StdAmrLoopVisitor.__init__(self, pathbook, manager, rundir)
  def solverPost(self, iter):
    sys.exit(1)

#------------------------------------------------------------------------

class DwrVisitor(StdAmrLoopVisitor):
  def __init__(self, pathbook, manager, rundir) :
    StdAmrLoopVisitor.__init__(self, pathbook, manager, rundir)
  def solverPost(self, iter):
    model = self.manager.parameters[u"Manager"][u"model"]._getvalue()
    modeladjoint = model + u"_Adjoint"
    self.manager.parameters[u"Manager"][u"model"] = modeladjoint
    loop = u'connectivity'
    cg = self.manager.runProject(loop=loop)
    if cg.return_code != 0 :
      print u'$$$ Solver() commande \"', loop, u'\" n\'a pas marché !!!'
      sys.exit(1)
    loop = u'linear'
    cg = self.manager.runProject(loop=loop)
    if cg.return_code != 0 :
      print u'$$$ Solver() commande \"', loop, u'\" n\'a pas marché !!!'
      sys.exit(1)
    loop = u'postprocess'
    cg = self.manager.runProject(loop=loop)
    if cg.return_code != 0 :
      print u'$$$ Solver() commande \"', loop, u'\" n\'a pas marché !!!'
      sys.exit(1)
    vtkfile = self.pathbook(u'project', self.rundir, modeladjoint + u'.vtk')
    vtkfile2 = self.pathbook(u'project', self.rundir, modeladjoint + u'.%06d.vtk' % (iter))
    if not vtkfile.exists() :
      print u'$$$ Solver() fichier \"', vtkfile, u'\" n\'existe pas !!!'
      sys.exit(1)
    vtkfile.move(vtkfile2)
    errorfile = self.pathbook(u'project', self.rundir, u'PostProcess' + u'.pylotsolution')
    errorfile2 = self.pathbook(u'project', self.rundir, modeladjoint + u'_PostProcess' + u'.pylotsolution')
    errorfile.move(errorfile2)

    self.manager.parameters[u"Manager"][u"model"] = model


#------------------------------------------------------------------------

class AmrLoop() :
  def __init__(self, firstmesh, param, pathbook, manager, rundir, visitor) :
    self.firstmesh = firstmesh
    self.param = param
    self.pathbook = pathbook
    self.manager = manager
    self.rundir = rundir
    self.visitor = visitor
    self.niter = int(param[u"AdaptiveAlgorithm"][u"niter"].value)
    # self.model = rundir.split('/')[-3]
    self.model = param[u"Problem"][u"model"].value
    self.solverloop = param[u"AdaptiveAlgorithm"][u"solverloop"].value
    marker = param[u"AdaptiveAlgorithm"][u"marker"].value
    marking_percentage = param[u"AdaptiveAlgorithm"][u"marking_percentage"].value
    gamma = param[u"AdaptiveAlgorithm"][u"gamma"].value
    self.filename = self.visitor.getResultsFileName(marker, marking_percentage, gamma)
    filehandle = open(self.filename, u'w')
    filehandle.close()
    self.meshname = self.pathbook(u'project', self.rundir, u'start')
    self.nextmeshname = self.pathbook(u'project', self.rundir, u'start2')
    if param.hasDatum(u"AdaptiveAlgorithm", u"quadtotri"):
      self.quadtotri = param[u"AdaptiveAlgorithm"][u"quadtotri"] == u"yes"
    else:
      self.quadtotri = False
    if self.quadtotri:
      self.transformedmeshname = self.pathbook(u'project', self.rundir, u'transformed')
    else:
      self.transformedmeshname = self.meshname
    self.datatype = u"binary"
    if self.param.hasDatum(u"Problem", u"datatype") :
      self.datatype = self.param[u"Problem"][u"datatype"].value
    self.manager.parameters[u"IoManager"][u"datatype"] = self.datatype
    self.manager.parameters[u"IoManager"][u"rundir"] = self.rundir
    self.manager.parameters[u"IoManager"][u"inmesh"] = self.transformedmeshname.split(u'/')[-1]
    self.manager.parameters[u"Manager"][u"application"] = self.param[u"Problem"][u"application"].value
    self.manager.parameters[u"Manager"][u"model"] = self.model
    if param.hasDatum(u"Problem", u"exactsolution"):
      self.manager.parameters[u"Application"][u"exactsolution"] = param[u"Problem"][u"exactsolution"].value
    self.manager.parameters[u"AdaptiveAlgorithm"][u"marker"] = self.param[u"AdaptiveAlgorithm"][u"marker"].value
    self.manager.parameters[u"AdaptiveAlgorithm"][u"marking_percentage"] = self.param[u"AdaptiveAlgorithm"][u"marking_percentage"].value
    self.manager.parameters[u"AdaptiveAlgorithm"][u"gamma"] = self.param[u"AdaptiveAlgorithm"][u"gamma"].value
    self.manager.parameters.saveParameters()


  def Refiner(self) :
    if self.iter == 0:
      nprerefine = self.param[u"AdaptiveAlgorithm"][u"nprerefine"].value
      executable = self.pathbook(u'bin', u'MeshReaders', u'ReaderQuad')
      args = [self.firstmesh, self.meshname, self.datatype]
      cg = self.manager.runCmd(cmd=executable, args=args)
      if cg.return_code != 0 :
        print u'$$$  commande \"', executable, u'\" n\'a pas marché !!!'
        sys.exit(1)
      executable = self.pathbook(u'bin', u'MeshTools', u'pyLotMeshUniformRefiner')
      args = [self.meshname, self.meshname, unicode(nprerefine)]
      cg = self.manager.runCmd(cmd=executable, args=args)
      if cg.return_code != 0 :
        print u'$$$  commande \"', executable, u'\" n\'a pas marché !!!'
        sys.exit(1)
    else :
      executable = self.pathbook(u'bin', u'MeshTools', u'QuadMeshRefiner')
      args = [self.meshname, self.nextmeshname, self.datatype]
      cg = self.manager.runCmd(cmd=executable, args=args)
      if cg.return_code != 0 :
        print u'$$$ Refiner() commande \"', executable, u'\" n\'a pas marché !!!'
        sys.exit(1)
      mesh = self.meshname + u'.pylotmesh'
      nextmesh = self.nextmeshname + u'.pylotmesh'
      mesh.rmtree()
      nextmesh.copytree(mesh)
      adaptmesh = self.pathbook(u'project', self.rundir, u'start.pylotmesh', u'start2.pylotmeshadaptive')
      adaptmesh2 = self.pathbook(u'project', self.rundir, u'start.pylotmesh', u'start.pylotmeshadaptive')
      adaptmesh.move(adaptmesh2)
    if self.quadtotri:
      executable = self.pathbook(u'bin', u'MeshTools', u'QuadMeshToTriangleMesh')
      args = [self.meshname, self.transformedmeshname, u"crisscross", self.datatype]
      cg = self.manager.runCmd(cmd=executable, args=args)
      if cg.return_code != 0 :
        print u'$$$ Refiner() commande \"', executable, u'\" n\'a pas marché !!!'
        sys.exit(1)

    executable = self.pathbook(u'bin', u'MeshTools', u'CreateNormals')
    args = [self.transformedmeshname]
    cg = self.manager.runCmd(cmd=executable, args=args)
    if cg.return_code != 0 :
      print u'$$$ Refiner() commande \"', executable, u'\" n\'a pas marché !!!'
      sys.exit(1)


  def Solver(self) :
    loop = u'connectivity'
    cg = self.manager.runProject(loop=loop)
    if cg.return_code != 0 :
      print u'$$$ Solver() commande \"', loop, u'\" n\'a pas marché !!!'
      sys.exit(1)
    loop = self.solverloop
    cg = self.manager.runProject(loop=loop)
    if cg.return_code != 0 :
      print u'$$$ Solver() commande \"', loop, u'\" n\'a pas marché !!!'
      sys.exit(1)
    vtkfile = self.pathbook(u'project', self.rundir, self.model + u'.vtk')
    vtkfile2 = self.pathbook(u'project', self.rundir, self.model + u'.%06d.vtk' % (self.iter))
    if not vtkfile.exists() :
      print u'$$$ Solver() fichier \"', vtkfile, u'\" n\'existe pas !!!'
      sys.exit(1)
    vtkfile.move(vtkfile2)
    loop = u'postprocess'
    cg = self.manager.runProject(loop=loop)
    if cg.return_code != 0 :
      print u'$$$ Solver() commande \"', loop, u'\" n\'a pas marché !!!'
      sys.exit(1)
    errorfile = self.pathbook(u'project', self.rundir, u'PostProcess' + u'.pylotsolution')
    errorfile2 = self.pathbook(u'project', self.rundir, self.model + u'_PostProcess' + u'.pylotsolution')
    errorfile.move(errorfile2)
    self.visitor.solverPost(self.iter)


  def Marker(self) :
    executable = self.pathbook(u'bin', u'MeshTools', u'EstimatorToMarkedCells')
    args = [self.pathbook(u'projectparam'), self.model]
    cg = self.manager.runCmd(cmd=executable, args=args)
    if cg.return_code != 0 :
      print u'$$$ Marker() commande \"', executable, u'\" n\'a pas marché !!!'
      sys.exit(1)

    if self.quadtotri:
      executable = self.pathbook(u'bin', u'MeshTools', u'MarkedCellsTriangleMeshToQuadMesh')
      args = [self.transformedmeshname, self.meshname]
      cg = self.manager.runCmd(cmd=executable, args=args)
      if cg.return_code != 0 :
        print u'$$$ Marker() commande \"', executable, u'\" n\'a pas marché !!!'
        sys.exit(1)

    meshinfo = MeshInfo(self.transformedmeshname)
    iterationinfo = IterationInfo(self.rundir, self.model)
    filehandle = open(self.filename, u'a')
    self.visitor.write(filehandle, self.iter, meshinfo, iterationinfo)
    filehandle.close()
    if meshinfo.ncells > float(self.param[u"AdaptiveAlgorithm"][u"nmax"].value) :
      sys.exit(99)

  def Run(self):
    for iter in range(self.niter) :
      self.iter = iter
      self.Refiner()
      self.Solver()
      self.Marker()
      if iter == self.niter - 1 : break

